package com.example.demo.service;

import com.example.demo.util.PoiExcelUtil;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.model.HasKey;
import us.codecraft.webmagic.pipeline.PageModelPipeline;
import us.codecraft.webmagic.utils.FilePersistentBase;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletResponse;
import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 *  注解模式获取数据处理
 * @author: wbx
 */
@Slf4j
public class IpSpiderServicePipeline extends FilePersistentBase implements PageModelPipeline<IpSpieder> {
    @Resource
    HttpServletResponse response;

    public IpSpiderServicePipeline() {
        setPath("/data/webmagic/");
    }

    public IpSpiderServicePipeline(String path) {
        setPath(path);
    }

    // 标志位，判断是否为首行
    private AtomicBoolean flag = new AtomicBoolean(true);

    @SneakyThrows
    @Override
    public void process(IpSpieder ipSpieder, Task task) {
        String path = this.path + PATH_SEPERATOR + task.getUUID() + PATH_SEPERATOR;

        BufferedWriter writer = null;
        // excel数据
        List<Map<String,String>> dataList =  new ArrayList<>();
        // excel导出路径
//        String excelFilePath = path + ((HasKey) ipSpieder).key() + ".xls";
        LinkedHashMap<String,String> title = new LinkedHashMap<>();
        title.put("ip","IP");
        title.put("port","PORT");
        title.put("anonymous","匿名度");
        title.put("type","类型");
        title.put("position","位置");
        title.put("responseSpeed","响应速度");
        title.put("validationTime","最后验证时间");

        try {
            String filename;
            if (ipSpieder instanceof HasKey) {
                filename = path + ((HasKey) ipSpieder).key() + ".txt";
            } else {
                filename = path + "ipAddress.txt";
            }
            writer = new BufferedWriter(new FileWriter(getFile(filename), true));
            if (flag.get()) {
                writer.write("IP\tPORT\t匿名度\t类型\t位置\t响应速度\t最后验证时间\r\n");
                flag.set(false);
            }

            List<String> ip = ipSpieder.getIp();
            List<String> port = ipSpieder.getPort();
            List<String> anonymous = ipSpieder.getAnonymous();
            List<String> type = ipSpieder.getType();
            List<String> position = ipSpieder.getPosition();
            List<String> responseSpeed = ipSpieder.getResponseSpeed();
            List<String> validationTime = ipSpieder.getValidationTime();
            for (int i = 0; i < ip.size(); i++) {
                Map<String,String> map = new HashMap<>();
                map.put("ip",ip.get(i));
                map.put("port",port.get(i));
                map.put("anonymous",anonymous.get(i));
                map.put("type",type.get(i));
                map.put("position",position.get(i));
                map.put("responseSpeed",responseSpeed.get(i));
                map.put("validationTime",validationTime.get(i));
                dataList.add(map);
                writer.write(ip.get(i) + "\t" + port.get(i) + "\t" + anonymous.get(i) + "\t" + type.get(i) + "\t"
                        + position.get(i) + "\t" + responseSpeed.get(i) + "\t" + validationTime.get(i) + "\r\n");
            }
            PoiExcelUtil.listToExcelModifyFileName(dataList, title, "数据信息", dataList.size() + 2, response, "测试",path + "数据表.xls");
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (writer != null) {
                try {
                    writer.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }


}
